#include "llvm/Transforms/MPI/MPI.h"
#include "llvm/Analysis/LibCallAliasAnalysis.h"
#include "llvm/Analysis/LibCallSemantics.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/Passes.h"
#include "llvm/Analysis/CFGPrinter.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Function.h"
#include "llvm/Support/CallSite.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Value.h"
#include "llvm/Pass.h"
#include "llvm/Instructions.h"
#include <algorithm>
#include <set>
#include <map>
using namespace llvm;


namespace {

  class MPILibCallInfo : public LibCallInfo {
  private:

    enum LocationIDs {
      StartLocID = -1,
      IBufLocID = 0,
      IDatatypeLocID,
      ICommLocID,
      IRequestLocID,
      WaitRequestLocID,
      WaitStatusLocID,
      EndLocID
    };

    static const LibCallFunctionInfo::LocationMRInfo IrecvLocDetails[];
    static const LibCallFunctionInfo IrecvFI;
    static const LibCallFunctionInfo::LocationMRInfo IsendLocDetails[];
    static const LibCallFunctionInfo IsendFI;
    static const LibCallFunctionInfo::LocationMRInfo WaitLocDetails[];
    static const LibCallFunctionInfo WaitFI;

    static const LibCallFunctionInfo FunctionInfoArray[];

    // TODO: I'm not sure how we are supposed to check locations
    // without AA, so I'm just *really* hacking one in for now.

    static LibCallLocationInfo::LocResult
    pointerLocationQuery(const Value *Loc, unsigned LocSize,
                         const Value *Ptr, unsigned Size) {
      assert(MPISupport::MPIAA && "AA not initialized yet?");
      AliasAnalysis *AA = MPISupport::MPIAA;
      // Check if they alias or not
      AliasAnalysis::AliasResult Res = AA->alias(Loc, LocSize, Ptr, Size);
      if(Res == AliasAnalysis::NoAlias)
        return LibCallLocationInfo::No;
      if(Res == AliasAnalysis::MustAlias)
        return LibCallLocationInfo::Yes;
      return LibCallLocationInfo::Unknown;
    }

    static LibCallLocationInfo::LocResult
    isLocationIBuf(ImmutableCallSite CS, const Value *Ptr, unsigned Size);

    static LibCallLocationInfo::LocResult
    isLocationIDatatype(ImmutableCallSite CS, const Value *Ptr, unsigned Size);

    static LibCallLocationInfo::LocResult
    isLocationIComm(ImmutableCallSite CS, const Value *Ptr, unsigned Size);

    static LibCallLocationInfo::LocResult
    isLocationIRequest(ImmutableCallSite CS, const Value *Ptr, unsigned Size);

    static LibCallLocationInfo::LocResult
    isLocationWaitRequest(ImmutableCallSite CS, const Value *Ptr, unsigned Size);

    static LibCallLocationInfo::LocResult
    isLocationWaitStatus(ImmutableCallSite CS, const Value *Ptr, unsigned Size);

    // Our local copy of Locations. LibCallInfo asks for it when needed.
    LibCallLocationInfo Locations[EndLocID];
  public:

    MPILibCallInfo() {
      Locations[IBufLocID].isLocation = isLocationIBuf;
      Locations[IDatatypeLocID].isLocation = isLocationIDatatype;
      Locations[ICommLocID].isLocation = isLocationIComm;
      Locations[IRequestLocID].isLocation = isLocationIRequest;
      Locations[WaitRequestLocID].isLocation = isLocationWaitRequest;
      Locations[WaitStatusLocID].isLocation = isLocationWaitStatus;
    }
    virtual ~MPILibCallInfo() {}

    virtual unsigned getLocationInfo(const LibCallLocationInfo *&Array) const;
    virtual const LibCallFunctionInfo *getFunctionInfoArray() const;
  };

} // end of namespace

unsigned
MPILibCallInfo::getLocationInfo(const LibCallLocationInfo *&Array) const {
  Array = Locations;
  return EndLocID; // Number of Locations
}

const LibCallFunctionInfo *MPILibCallInfo::getFunctionInfoArray() const {
  return MPILibCallInfo::FunctionInfoArray;
}

// TODO: How to invalidate buf?
const LibCallFunctionInfo::LocationMRInfo MPILibCallInfo::IrecvLocDetails[] = {
    { IBufLocID, AliasAnalysis::Mod },
    { IDatatypeLocID, AliasAnalysis::Ref },
    { ICommLocID, AliasAnalysis::Ref },
    { IRequestLocID, AliasAnalysis::Mod },
    { ~0U, AliasAnalysis::ModRef } // array end
};

const LibCallFunctionInfo MPILibCallInfo::IrecvFI = {
    "MPI_Irecv",            // Name
    AliasAnalysis::ModRef,  // UniversalBehavior
    LibCallFunctionInfo::DoesOnly, // DetailsType
    IrecvLocDetails
};

const LibCallFunctionInfo::LocationMRInfo MPILibCallInfo::IsendLocDetails[] = {
    { IBufLocID, AliasAnalysis::Ref },
    { IDatatypeLocID, AliasAnalysis::Ref },
    { ICommLocID, AliasAnalysis::Ref },
    { IRequestLocID, AliasAnalysis::Mod },
    { ~0U, AliasAnalysis::ModRef } // array end
};

const LibCallFunctionInfo MPILibCallInfo::IsendFI = {
    "MPI_Isend",            // Name
    AliasAnalysis::ModRef,  // UniversalBehavior
    LibCallFunctionInfo::DoesOnly, // DetailsType
    IsendLocDetails
};

const LibCallFunctionInfo::LocationMRInfo MPILibCallInfo::WaitLocDetails[] = {
    { WaitRequestLocID, AliasAnalysis::ModRef },
    { WaitStatusLocID, AliasAnalysis::Mod },
    { ~0U, AliasAnalysis::ModRef } // array end
};

const LibCallFunctionInfo MPILibCallInfo::WaitFI = {
    "MPI_Wait",            // Name
    AliasAnalysis::ModRef,  // UniversalBehavior
    LibCallFunctionInfo::DoesOnly, // DetailsType
    WaitLocDetails
};

const LibCallFunctionInfo MPILibCallInfo::FunctionInfoArray[] = {
    MPILibCallInfo::IrecvFI,
    MPILibCallInfo::IsendFI,
    MPILibCallInfo::WaitFI,
    { 0, AliasAnalysis::ModRef, LibCallFunctionInfo::DoesOnly, 0 } // end of array
};

// For MPI LibCalls we are mainly concerned with parameter locations.
// For a given CS and Ptr, we must tell if Ptr aliases what is being
// passed to the specified parameter location. TODO: verify correctness?

LibCallLocationInfo::LocResult
MPILibCallInfo::isLocationIBuf(ImmutableCallSite CS, const Value *Ptr,
                                        unsigned Size) {
  const Value *Buf = MPISupport::getPointerArg(CS, MPIArgs::IBuf);
  return MPILibCallInfo::pointerLocationQuery(Buf, ~0U, Ptr, Size);
}

LibCallLocationInfo::LocResult
MPILibCallInfo::isLocationIDatatype(ImmutableCallSite CS, const Value *Ptr,
                                             unsigned Size) {
  // The MPI Standard defines MPI_Datatype as an opaque object, so we need
  // only verify that the underlying object and the parameter are the same
  const Value *Datatype = CS.getArgument(MPIArgs::IDatatype)->getUnderlyingObject();
  if(Datatype == Ptr->getUnderlyingObject())
    return LibCallLocationInfo::Yes;

  // TODO: what if memcpy or anything similar?
  // Let's be conservative for now.
  // XXX: Scratch that, the current behavior of LibCallAA is too conservative,
  // It is not properly using the DoesOnly case, so we force a No here for now.
  return LibCallLocationInfo::No;
}

LibCallLocationInfo::LocResult
MPILibCallInfo::isLocationIComm(ImmutableCallSite CS, const Value *Ptr,
                                         unsigned Size) {
  // The MPI Standard defines MPI_Comm as an opaque object, so we need
  // only verify that the underlying object and the parameter are the same
  const Value *Comm = CS.getArgument(MPIArgs::IComm)->getUnderlyingObject();
  if(Comm == Ptr->getUnderlyingObject())
    return LibCallLocationInfo::Yes;

  // TODO: what if memcpy or anything similar?
  // Let's be conservative for now
  // XXX: Scratch that, the current behavior of LibCallAA is too conservative,
  // It is not properly using the DoesOnly case, so we force a No here for now.
  return LibCallLocationInfo::No;
}

LibCallLocationInfo::LocResult
MPILibCallInfo::isLocationIRequest(ImmutableCallSite CS, const Value *Ptr,
                                            unsigned Size) {
  const Value *Request = MPISupport::getPointerArg(CS, MPIArgs::IRequest);
  return MPILibCallInfo::pointerLocationQuery(Request, ~0U, Ptr, Size);
}

LibCallLocationInfo::LocResult
MPILibCallInfo::isLocationWaitRequest(ImmutableCallSite CS, const Value *Ptr,
                                               unsigned Size) {
  const Value *Request = MPISupport::getPointerArg(CS, MPIArgs::WaitRequest);
  return MPILibCallInfo::pointerLocationQuery(Request, ~0U, Ptr, Size);
}

LibCallLocationInfo::LocResult
MPILibCallInfo::isLocationWaitStatus(ImmutableCallSite CS, const Value *Ptr,
                                              unsigned Size) {
  const Value *Status = MPISupport::getPointerArg(CS, MPIArgs::WaitStatus);
  return MPILibCallInfo::pointerLocationQuery(Status, ~0U, Ptr, Size);
}

namespace {

 class MPIAliasAnalysis : public LibCallAliasAnalysis {
  private:
    // TODO: switch to LLVM data structures
    std::map<Value*, std::vector<Value*> > RequestToBuf;

    static bool verifyValidFType(Function *F) {
      if(F->getName() == "MPI_Irecv" || F->getName() == "MPI_Isend")
        return MPISupport::verifyIrecvSignature(F->getFunctionType(), &F->getContext());
      else if(F->getName() == "MPI_Wait")
        return MPISupport::verifyWaitSignature(F->getFunctionType(), &F->getContext());

      assert(0 && "Unknown Function.");
    }

  public:
    static char ID; // Class identification

    explicit MPIAliasAnalysis()
      : LibCallAliasAnalysis(ID, 0){
      LCI = new MPILibCallInfo();
    }
    ~MPIAliasAnalysis() {}

    // We do the analysis of effects hidden by the library in these methods
    ModRefResult getModRefInfo(CallSite CS, Value *P, unsigned Size);
    ModRefResult getModRefInfo(CallSite CS1, CallSite CS2);

    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.addRequiredID(MPIAnnotateFunctionsID);
      AliasAnalysis::getAnalysisUsage(AU);
      AU.setPreservesAll();
    }

    virtual bool runOnFunction(Function &F);
  };

} // end of namespace

// Register this pass...
char MPIAliasAnalysis::ID = 0;
static RegisterPass<MPIAliasAnalysis> X("mpi-aa", "MPI Alias Analysis", true, true);
//INITIALIZE_PASS(MPIAliasAnalysis,"mpi-aa", "MPI Alias Analysis", true, true);

// Declare that we implement the AliasAnalysis interface
static RegisterAnalysisGroup<AliasAnalysis> Y(X);

AliasAnalysis *MPISupport::MPIAA = 0;

const PassInfo *const MPIAliasAnalysisID = &X;
FunctionPass *createMPIAliasAnalysisPass() {
  MPIAliasAnalysis *MPIAAPass = new MPIAliasAnalysis();
  MPISupport::MPIAA = MPIAAPass;
  return MPIAAPass;
}

bool MPIAliasAnalysis::runOnFunction(Function &F) {
  InitializeAliasAnalysis(this);

  for (inst_iterator II = inst_begin(F), E = inst_end(F); II != E; ++II) {
    if (!isa<CallInst>(*II) && !isa<InvokeInst>(*II))
      continue;

    CallSite CS = CallSite::get(&*II);
    Function *F = CS.getCalledFunction();

    if(!LCI->getFunctionInfo(F))
      continue;

    StringRef FName = F->getName();
    if(FName == "MPI_Irecv" || FName == "MPI_Isend") {
      // Verify MPI_Irecv and MPI_Isend signature TODO: Datatype and Comm?
      // They both have the same signature, so we can call either one.
      if(!MPISupport::verifyIrecvSignature(F->getFunctionType(), &F->getContext()))
        continue;

      // Irecv and Isend analysis, FIXME: This definitely broken, figure out what we need
      Value *Request = CS.getArgument(MPIArgs::IRequest)->getUnderlyingObject();
      RequestToBuf[Request].push_back(CS.getArgument(MPIArgs::IBuf)); // whichbuf[request] = buf
    }

    // TODO: Map Irecv/Isend to Waits ?
  }

  return false;
}

AliasAnalysis::ModRefResult
MPIAliasAnalysis::getModRefInfo(CallSite CS, Value *P, unsigned Size) {
  Function *F = CS.getCalledFunction();
  StringRef FName = F->getName();

  if(FName == "MPI_Wait") {
    if(MPISupport::verifyWaitSignature(F->getFunctionType(), &F->getContext())) {
      // MPI_Wait Mod's the buffer used on the original non-blocking call,
      // but we need to say it also Ref's so that redefining it in between
      // both calls is invalid.
      // TODO: Do we also need to check any alias for the requests?
      std::vector<Value*> *Bufs = &RequestToBuf[CS.getArgument(MPIArgs::WaitRequest)];
      assert(Bufs && "Unmapped Request argument for MPI_Wait.");

      for(std::vector<Value*>::const_iterator BI = Bufs->begin(), E = Bufs->end(); BI != E; ++BI) {
        Value *Buf = *BI;
        // TODO: Try to get size of pointer, see PointerTracking::computeAllocationCount
//        if(AA->alias(Buf, ~0U, P, Size))
//          // Alias returned MayAlias or MustAlias
//          return ModRef;
      }
    }
  }

  return LibCallAliasAnalysis::getModRefInfo(CS, P, Size);
}

AliasAnalysis::ModRefResult
MPIAliasAnalysis::getModRefInfo(CallSite CS1, CallSite CS2) {
  Function *F1 = CS1.getCalledFunction();
  StringRef F1Name = F1->getName();

  Function *F2 = CS2.getCalledFunction();
  StringRef F2Name = F2->getName();

  // Make sure we know about CS1's called function
  if(!LCI->getFunctionInfo(F1) || !verifyValidFType(F1))
    return LibCallAliasAnalysis::getModRefInfo(CS1,CS2);

  // If we don't know about the other function, treat each arguments separately
  if(!LCI->getFunctionInfo(F2) || !verifyValidFType(F2)) {
    ModRefResult MRInfo = NoModRef;
    for(CallSite::arg_iterator AI = CS2.arg_begin(), AE = CS2.arg_end(); AI != AE; ++AI) {
      Value* V = dyn_cast<Value>(AI);
      if(isa<PointerType>(V->getType()))
        MRInfo = (ModRefResult)(MRInfo | LibCallAliasAnalysis::getModRefInfo(CS1, V, ~0U));
    }
    return MRInfo;
  }

  // We have two known MPI function calls
  if(F1Name == "MPI_Irecv" || F1Name == "MPI_Isend") {
    // XXX: I'm pretty sure we can treat them as the same, but double-check later
    if(F2Name == "MPI_Isend" || F2Name == "MPI_Irecv") {
      ModRefResult MRInfo = NoModRef; // we start optimistic, and exit early if we find ModRef

      // TODO: Try to get size of pointer, see PointerTracking::computeAllocationCount
      // We can also try using count here, but I'm not sure how LLVM sizes work
//      if(AA->alias(CS1.getArgument(MPIArgs::IBuf), ~0U,
//          CS2.getArgument(MPIArgs::IBuf), ~0U))
//        return ModRef; // It's invalid to use the same buf before completion

//      if(AA->alias(CS1.getArgument(MPIArgs::IRequest), ~0U,
//          CS2.getArgument(MPIArgs::IRequest), ~0U))
//        MRInfo = (ModRefResult)(MRInfo | Mod);

      // Now we try to match (src, tag, comm) if they're the same
      // XXX: Note that we don't care about enforcing ordering.
      if(F1Name == F2Name) {
        // TODO: Check Communicators. We can probably assume that a simple
        // alias check is enough if comm_dup hasn't been called.
        bool SameComm = true;

        // TODO: We can probably do a better job trying to find int
        // constants through PHI nodes and possible feeds to loads
        bool SameTag = true;
        ConstantInt *FirstTag =
          dyn_cast<ConstantInt>(CS1.getArgument(MPIArgs::ITag));
        ConstantInt *SecondTag =
          dyn_cast<ConstantInt>(CS2.getArgument(MPIArgs::ITag));

        if(FirstTag && SecondTag) {
          if(FirstTag->getSExtValue() != MPIConstants::ANY_TAG &&
              SecondTag->getSExtValue() != MPIConstants::ANY_TAG &&
              FirstTag->getSExtValue() != SecondTag->getSExtValue())
            SameTag = false;
        }

        // We can use the logic for Irecv on Isend's (it will just do dest != dest)
        bool SameSource = true;
        ConstantInt *FirstSource =
          dyn_cast<ConstantInt>(CS1.getArgument(MPIArgs::ISource));
        ConstantInt *SecondSource =
          dyn_cast<ConstantInt>(CS2.getArgument(MPIArgs::ISource));

        if(FirstSource && SecondSource) {
          /* TODO: Check the Standard to verify this is true:
          if(FirstSource->getSExtValue() == MPIConstants::PROC_NULL ||
              SecondSource->getSExtValue() == MPIConstants::PROC_NULL)
            return NoModRef;
          */

          if(FirstSource->getSExtValue() != MPIConstants::ANY_SOURCE &&
              SecondSource->getSExtValue() != MPIConstants::ANY_SOURCE &&
              FirstSource->getSExtValue() != SecondSource->getSExtValue())
            SameSource = false;
        }

        if(SameComm && SameTag && SameSource)
          return ModRef; // They read the same message
      }

      // We don't care about interleaving Isend and Irecv if they don't alias
      return MRInfo;
    }

    if(F2Name == "MPI_Wait") {
      // TODO: We could probably get the size of MPI_Request from the library
//      if(!AA->alias(CS1.getArgument(MPIArgs::IRequest), ~0U,
//                    CS2.getArgument(MPIArgs::WaitRequest), ~0U))
//        return NoModRef;
    }
  }

  if(F1Name == "MPI_Wait") {
    if(F2Name == "MPI_Wait") {
//      if(AA->alias(CS1.getArgument(MPIArgs::WaitRequest), ~0U,
//                    CS2.getArgument(MPIArgs::WaitRequest), ~0U))
//        return ModRef;
//
//      if(AA->alias(CS1.getArgument(MPIArgs::WaitStatus), ~0U,
//                    CS2.getArgument(MPIArgs::WaitStatus), ~0U))
//        return ModRef;

      return NoModRef;
    }

    if(F2Name == "MPI_Irecv" || F2Name == "MPI_Isend") {
//      if(AA->alias(CS1.getArgument(MPIArgs::WaitRequest), ~0U,
//                    CS2.getArgument(MPIArgs::IRequest), ~0U))
//        return ModRef;

      return NoModRef;
    }
  }

  // Chain to base classes
  return LibCallAliasAnalysis::getModRefInfo(CS1,CS2);
}
